# code.py

from rembg import remove, new_session
from PIL import Image
import os

def remove_background(input_path, output_path):
    # 读取输入图像
    input_img = Image.open(input_path)

    # 移除背景（自动生成透明背景）

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # session = new_session(model_name='u2net_human_seg')
    session = new_session(model_name='isnet-general-use')
    # session = new_session(model_name='silueta')
    output_img = remove(input_img, session=session)

    # 保存透明背景图像（PNG格式支持透明通道）
    # output_img.save(white_path)
    # 创建白色背景
    white_bg = Image.new("RGB", output_img.size, (255, 255, 255))

    # 将透明图像叠加到白色背景上
    white_bg.paste(output_img, mask=output_img.split()[3])  # 使用alpha通道作为蒙版

    # 保存结果
    white_bg.save(output_path)

# 使用示例
remove_background("../images/2025-07-01_140049_168.png", "method1/ggg.png")